using System;
using System.Collections.Generic;
using System.Text;
using Blimpbots.Recognition.Algorithms;
using System.Drawing;

namespace Blimpbots.Recognition
{
    public class WarpModel
    {
        public WarpModel(ShuffledList<KeypointMatch> inliers, IList<KeypointMatch> allMatches, double p_up)
        {
            _allMatches = allMatches;
            _scale = 0; _angle = 0;
            _inliers = inliers;
            P_up = p_up;

            EstimateParameters(100); //TODO: only call this when we like the model
        }

        private IList<KeypointMatch> _allMatches;
        private ShuffledList<KeypointMatch> _inliers;
        private double _scale;
        private double _surfScaleAverage;
        private double _angle;
        private double _angleStDev;

        public double P_up; // Probability of "up" orientation

        public IList<KeypointMatch> AllMatches { get { return _allMatches; } }
 
        public ShuffledList<KeypointMatch> Inliers { get { return _inliers; } }
        public double Scale { get {
            if (UseSurfScale) return _surfScaleAverage;
            else
            return (_surfScaleAverage * 2 + _scale) / 3.0; 
        } }

        public bool UseSurfScale = false;

        // Deviation between the model- and keypoint- derived scales
        public double ScaleDeviation { get { return Math.Abs(_surfScaleAverage -_scale); } }

        public double Angle { get { return _angle; } }

        public double AngleSdDev { get { return _angleStDev; } }

        /// <summary>
        /// Robustly estimate the scale of a detected object.
        /// </summary>
        /// <param name="sampleCount">Number of samples to pull</param>
        /// <returns>Scale estimate</returns>
        private void EstimateParameters(int sampleCount)
        {
            /* Scale Estimate Based on Keypoints Directly */
            double[] samples = new double[Inliers.Count];
            for (int i = 0; i < Inliers.Count; i++)
            {
                samples[i] = Inliers[i].dataScale / Inliers[i].modelScale;
            }

            _surfScaleAverage = ArrayUtils.Average(samples, 0, Inliers.Count);

            double angSqErr = 0;

            double phi, beta, scale;

            samples = new double[sampleCount];
            SimpleAngleAverage angleAverage = new SimpleAngleAverage();
            int sample_idx = 0;

            do // allow random inlier pairs to vote for the scale
            {
                _inliers.Shuffle(1);
                for (int i = 0; i < _inliers.Count - 1; i = i + 1)
                {
                    ComputeTransform(_inliers[i], _inliers[i+1], out phi, out beta, out scale);

                    double aD = Math.Abs(Angles.Difference(beta, Angles.Normalize(_inliers[i].beta)));
                    aD += Math.Abs(Angles.Difference(beta, Angles.Normalize(_inliers[i+1].beta)));
                    double betaN = Angles.Normalize(beta);
                    // now the angle goes from 0...2Pi
                    //int bin = (int)((betaN * angleBins.Length) / (Math.PI * 2));
                    //angleBins[bin] += 1.0 / ((aD*aD) + 0.5); 

                    samples[sample_idx++] = scale;
                    angleAverage.Update(beta);

                    angSqErr += Math.Pow((beta - angleAverage.Angle()), 2);

                    if (sample_idx == sampleCount) break;
                }

            } while (sample_idx < sampleCount);

            _angleStDev = Math.Sqrt(angSqErr / sampleCount);

            //double bestBin = 0; int bestBinIdx = 0;
            //for (int i = 0; i < angleBins.Length; i++)
            //{
            //    if (angleBins[i] > bestBin)
            //    {
            //        bestBin = angleBins[i];
            //        bestBinIdx = i;
            //    }
            //}
            //_angle = (bestBinIdx * Math.PI * 2) / angleBins.Length;

            Array.Sort(samples);

            // Sum over the middle half of the data points
            double scale_sum = 0;
            for (int i = sampleCount / 4; i < sampleCount * 3 / 4; i++)
            { scale_sum += samples[i]; }

            // Return the average
            _scale = scale_sum / (sampleCount / 2);
            _angle = angleAverage.Angle();

        }

        public List<surfInterop.SurfKeypoint> FilterKeypoints(IList<surfInterop.SurfKeypoint> kpts, float padding)
        {
            List<surfInterop.SurfKeypoint> results = new List<surfInterop.SurfKeypoint>();

            RectangleF bound = Bound();
            bound.Inflate(padding, padding);

            foreach (surfInterop.SurfKeypoint kpt in kpts)
            {
                PointF pt = new PointF((float)kpt.x, (float)kpt.y);
                if (bound.Contains(pt))
                    results.Add(kpt);
            }

            return results;
        }

        public RectangleF Bound()
        {
            double min_x = double.MaxValue, min_y = double.MaxValue, max_x =0, max_y = 0;

            foreach (KeypointMatch kpm in _inliers)
            {
                min_x = Math.Min(kpm.dataCoord.x, min_x);
                min_y = Math.Min(kpm.dataCoord.y, min_y);
                max_x = Math.Max(kpm.dataCoord.x, max_x);
                max_y = Math.Max(kpm.dataCoord.y, max_y);
            }

            return RectangleF.FromLTRB((float)min_x, (float)min_y, (float)max_x, (float)max_y);
        }

        public static void ComputeTransform(KeypointMatch a, KeypointMatch b, out double phi, out double beta, out double scale)
        {
            scale = a.dataCoord.Distance(b.dataCoord) /
                    a.modelCoord.Distance(b.modelCoord);

            phi = Math.Atan2(b.modelCoord.y - a.modelCoord.y,
                             b.modelCoord.x - a.modelCoord.x);

            double theta
                = Math.Atan2(b.dataCoord.y - a.dataCoord.y,
                             b.dataCoord.x - a.dataCoord.x);

            beta = theta - phi;

        }

        public static Point2D PredictUpright(Point2D modelCoord, KeypointMatch a, double beta, double phi, double scale)
        {
            // Angle of this point from the current center
            double alpha = Math.Atan2(modelCoord.y - a.modelCoord.y,
                                      modelCoord.x - a.modelCoord.x);

            // Offset by the data-relative orientation and scale
            double a1 = alpha + beta;

            double d = scale * modelCoord.Distance(a.modelCoord);

            return new Point2D(
                a.dataCoord.x + d * Math.Cos(a1),
                a.dataCoord.y + d * Math.Sin(a1));
        }

        public static Point2D[] Predict(Point2D modelCoord, KeypointMatch a, double beta, double phi, double scale)
        {
            // Angle of this point from the current center
            double alpha = Math.Atan2(modelCoord.y - a.modelCoord.y,
                                      modelCoord.x - a.modelCoord.x);

            // Offset by the data-relative orientation and scale
            double a1 = alpha + beta;
            double a2 = 2 * phi - alpha + beta;

            double d = scale * modelCoord.Distance(a.modelCoord);

            Point2D[] ret = new Point2D[2];

            ret[0] = new Point2D(
                a.dataCoord.x + d * Math.Cos(a1),
                a.dataCoord.y + d * Math.Sin(a1));

            ret[1] = new Point2D(
                a.dataCoord.x + d * Math.Cos(a2),
                a.dataCoord.y + d * Math.Sin(a2));

            return ret;
        }

        public Point2D PredictPoint(Point2D modelPt)
        {
            return PredictPoint(modelPt, 50);
        }

        public Point2D PredictPoint(Point2D modelPt, int sampleCount)
        {
            double s = Scale;
            ShuffledList<KeypointMatch> inliers = new ShuffledList<KeypointMatch>(Inliers);
            Point2D avg = new Point2D();
            double wtSum = 0;

            int i = 0;
            for (int samples = 0; samples < sampleCount; samples++)
            {
                if ((i + 1) == inliers.Count)
                {
                    inliers.Shuffle(1);
                    i = 0;
                }

                double phi, beta, scale;
                WarpModel.ComputeTransform(inliers[i], inliers[i + 1], out phi, out beta, out scale);
                Point2D[] preds = Predict(modelPt, inliers[i], beta, phi, scale);

                // Choose the upright or downfacing according to the model's overall choice
                Point2D p = (P_up > 0.5) ? preds[0] : preds[1];

                // If the transform differs largely, it may not be valid, so weaken it
                double scaleDivergence = Math.Abs(scale - s);
                double angDivergence = Math.Abs(Angles.Difference(beta, _angle));
                
                double betaDivergence = Math.Abs(Angles.Difference(beta, inliers[i].beta));
                betaDivergence *= Math.Abs(Angles.Difference(beta, inliers[i+1].beta));

                betaDivergence *= 30;
                scaleDivergence *= scaleDivergence * 50;
                angDivergence *= angDivergence * 10;

                // Ensure the scale's within an acceptable range
                if (scale > RansacWarpFit.MinScaleAllowed && scale < RansacWarpFit.MaxScaleAllowed)
                {
                    double wt = 1.0/ (1.0 + Math.Pow(inliers[i].modelCoord.Distance(modelPt)/5, 2) + scaleDivergence);
                    avg.x += wt * p.x; avg.y += wt * p.y; wtSum += wt;
                }

                i++;
            }

            if (wtSum == 0) wtSum = 1.0; // TODO: do we need this?

            avg.x /= wtSum; avg.y /= wtSum;

            return avg;
        }
  
    }

    public sealed class RansacWarpFit
    {
        public const double MinScaleAllowed = 0.4;
        public const double MaxScaleAllowed = 3.3;


        private static readonly RansacWarpFit instance = new RansacWarpFit();
        public static RansacWarpFit Instance { get { return instance; } }

        private Graphics g;

        public Graphics DebugSurface { set { g = value; } }

        public WarpModel Run(IList<KeypointMatch> matches, int maxRuns)
        {
            ShuffledList<KeypointMatch> shuffled = new ShuffledList<KeypointMatch>(matches);

            WarpModel bestModel = default(WarpModel);
            int bestInlierCount = 0;

            for (int i = 0; i < maxRuns; i++)
            {
                shuffled.Shuffle(1);

                /* errThresh values:
                 * 
                 * 3.0 seems on the low side
                 * 6.0 seems on the high side
                 * 
                 */

                WarpModel m = WarpFitUpright(shuffled, /* errThresh = */ 5.0);
                if (m != null && m.Inliers.Count > bestInlierCount)
                { bestModel = m; bestInlierCount = m.Inliers.Count; }
            }

            return bestModel;
        }

        private WarpModel WarpFitUpright(IList<KeypointMatch> shuffledMatches, double errThresh)
        {
            SimpleAngleAverage angleAverage = new SimpleAngleAverage();

            // randomly initialize with a match pair
            // draw predicted ellipse for all other matches

            ShuffledList<KeypointMatch> sl = new ShuffledList<KeypointMatch>(shuffledMatches);
            ShuffledList<KeypointMatch> inliers = new ShuffledList<KeypointMatch>(sl.GetRange(0, 2)); // random selection
            sl.RemoveRange(0, 2);

            if (g != null) // Show the randomly drawn inliers
            {
                //g.DrawEllipse(Pens.Red, Visualization.PointRectangleF(inliers[0].dataCoord, (float)(errThresh * 2)));
                //g.DrawEllipse(Pens.Red, Visualization.PointRectangleF(inliers[1].dataCoord, (float)(errThresh * 2)));
            }

            int runs = 30;
            double phi, beta, scale;

            long uniquePairCount = 1;
            long pairsChecked = 0;

            while (runs > 0)
            {
                inliers.Shuffle(1);
                runs--;

                // Quit early if we're not adding pairs fast enough
                if (pairsChecked > uniquePairCount) break;

                pairsChecked++;

                // Find the mini-transform implied by a pair of keypoint matches
                WarpModel.ComputeTransform(inliers[0], inliers[1], out phi, out beta, out scale);

                // Kills extreme scales because they probably are incorrect
                if (scale < MinScaleAllowed || scale > MaxScaleAllowed) break;

                angleAverage.Update(beta);

                //double aD = Math.Abs(Angles.Difference(beta, inliers[0].beta));
                //aD += Math.Abs(Angles.Difference(beta, inliers[1].beta));

                //if (aD > 3.0) break; // Don't try when keypoints point the wrong way

                //double aD = Angles.Difference(beta, angleAverage.Angle());
                //if (aD > Math.PI/4 || aD < -Math.PI/4 ) break;

                // Check which other matches agree with the transform
                for (int i = 0; i < sl.Count; i++)
                {
                    angleAverage.Update(beta);

                    Point2D predicted = WarpModel.PredictUpright(sl[i].modelCoord, inliers[0], beta, phi, scale);
                    double err = sl[i].dataCoord.Distance(predicted);

                    if (err < errThresh)
                    {
                        //double aD = Math.Abs(Angles.Difference(beta, inliers[0].beta));
                        //aD += Math.Abs(Angles.Difference(beta, inliers[1].beta));

                        uniquePairCount += inliers.Count;

                        inliers.Add(sl[i]);
                        sl.RemoveAtSwap(i); i--;

                        if (g != null) // Draw new inliers
                        {
                            g.DrawEllipse(Pens.Blue, (float)(predicted.x - errThresh),
                                                     (float)(predicted.y - errThresh),
                                                     (float)errThresh * 2, (float)errThresh * 2);
                        }
                    }

                }

            } // end while

            WarpModel mdl = new WarpModel(inliers, shuffledMatches, 1.0);

            if (double.IsNaN(mdl.Scale)) return null; // failed to find params

            return mdl;
        }

        private WarpModel WarpFit(IList<KeypointMatch> shuffledMatches, double errThresh)
        {
            // randomly initialize with a match pair
            // draw predicted ellipse for all other matches

            ShuffledList<KeypointMatch> sl = new ShuffledList<KeypointMatch>(shuffledMatches);
            ShuffledList<KeypointMatch> inliers = new ShuffledList<KeypointMatch>(sl.GetRange(0, 2)); // random selection
            sl.RemoveRange(0, 2);

            if (g != null) // Show the randomly drawn inliers
            {
                //g.DrawEllipse(Pens.Red, Visualization.PointRectangleF(inliers[0].dataCoord, (float)(errThresh * 2)));
                //g.DrawEllipse(Pens.Red, Visualization.PointRectangleF(inliers[1].dataCoord, (float)(errThresh * 2)));
            }

            int runs = 10;
            double phi, beta, scale;

            int upVotes = 0;
            int downVotes = 0;

            while (runs > 0)
            {
                inliers.Shuffle(1);
                runs--;

                // Find the mini-transform implied by a pair of keypoint matches
                WarpModel.ComputeTransform(inliers[0], inliers[1], out phi, out beta, out scale);

                // Check which other matches agree with the transform
                for (int i = 0; i < sl.Count; i++)
                {
                    Point2D[] pred = WarpModel.Predict(sl[i].modelCoord, inliers[0], beta, phi, scale);
                    double err1 = sl[i].dataCoord.Distance(pred[0]);
                    double err2 = sl[i].dataCoord.Distance(pred[1]);

                    Point2D predicted = new Point2D(); double err;

                    // Select the best projection of the two choices

                    if (err1 < err2)
                    {
                        predicted = pred[0]; err = err1;
                    }
                    else
                    {
                        predicted = pred[1]; err = err2;
                    }

                    if (err < errThresh)
                    {
                        if (Math.Abs(err1 - err2) > 2.0)
                        {
                            if (err1 < err2) upVotes++;
                            else downVotes++;
                        }

                        inliers.Add(sl[i]);
                        sl.RemoveAtSwap(i); i--;

                        if (g != null) // Draw new inliers
                        {
                            g.DrawEllipse(Pens.Blue, (float)(predicted.x - errThresh),
                                                     (float)(predicted.y - errThresh),
                                                     (float)errThresh * 2, (float)errThresh * 2);
                        }
                    }

                }

            } // end while

            double p_up = (upVotes + downVotes == 0) ? 0.5 :
                (double)upVotes / (upVotes + downVotes);

            WarpModel mdl = new WarpModel(inliers, shuffledMatches, p_up);

            return mdl;
        }


    }
}
